from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

from Embed_model import BGE_Embed
from yunjian.model import get_model
from qw_model import qw_model

# 智谱大模型
# model = get_model()

# 千问大模型
model = qw_model()

embedding = BGE_Embed()

# 加载知识库
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 3, "score_threshold": 0.5})

# 检索出的文档
result = retriever.invoke('快手云剪多少钱')
print(result)

def doc_formatter(docs):
    # 使用集合去重，然后再合并成字符串
    unique_docs = set([doc.page_content for doc in docs])
    return '\n\n'.join(unique_docs)


system_prompt = (
    " 你是一个聪明又有亲和力的销售客服客服，擅长用轻松、亲切的语气解答用户关于公司业务、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。"
    "1. 当用户问题不清楚时，主动询问用户具体需求。"
    "2.有关云剪业务，有抖音平台、快手平台；有关直发或者实探，有抖音平台、小红书平台，所以要问清楚用户需求"
    "3.抖音平台是按达人级别划分，其他平台是粉丝量划分（根据用户选择的平台对应问哪一个，选抖音平台就是问'需要几级达人'，选其他平台就是'有粉丝量要求吗’）"
    "4. 价格问题严格按照文档给出，不可自行编造。"
    "5. 当没有用户所说的地区，就先按该地区所属的省份算价格，如果所属省份都没有，说明该用户说的地区只能做全国不限地区的，就按全国不限IP的价格算。"
    "6. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，避免用'您'。"
    "7. 输出答案应简洁明了，字数控制在15字左右。"
    "\n\n"
    "{context}"
)

# 创建提示词模板
prompt = ChatPromptTemplate.from_messages([
    ('system', system_prompt),

    ('user', '{input}')
])

chain = (
    # context是上下文，传的检索器结果，input是用户问题
    {'context': retriever | doc_formatter, 'input': RunnablePassthrough()}
    | prompt
    | model
    | (lambda x: x.content)
    | StrOutputParser()
)

res = chain.invoke('快手云剪多少钱')
print(res)
